import pandas as pd
import requests
import time
import os
import re
from urllib.parse import quote

def search_pubchem(input_file, output_file):
    """
    Search PubChem for chemical information using CAS numbers or English names
    and update an Excel file with the results.
    
    Args:
        input_file (str): Path to the input Excel file
        output_file (str): Path to save the updated Excel file
    """
    # Load Excel file
    if not os.path.exists(input_file):
        print(f"Error: Input file '{input_file}' not found")
        return
    
    try:
        df = pd.read_excel(input_file)
        print(f"Loaded file with {len(df)} rows")
    except Exception as e:
        print(f"Error loading file: {e}")
        return
    
    # Check if necessary columns exist
    required_columns = ["CAS号", "英文名", "CID号", "SMILES"]
    for col in required_columns:
        if col not in df.columns:
            df[col] = ""  # Create column if it doesn't exist
    
    # Track progress
    total_rows = len(df)
    rows_updated = 0
    rows_skipped = 0
    rows_failed = 0
    
    print("\n" + "="*50)
    print("STARTING PUBCHEM API SEARCH")
    print("="*50 + "\n")
    
    # Process each row
    for index, row in df.iterrows():
        # Get compound identifiers
        cas = row["CAS号"]
        english_name = row["英文名"]
        cid = row["CID号"]
        smiles = row["SMILES"]
        
        # Print compound being processed
        compound_id = english_name if not pd.isna(english_name) and english_name != "" else cas
        compound_id = compound_id if not pd.isna(compound_id) and compound_id != "" else f"Row {index+1}"
        print(f"\n[{index+1}/{total_rows}] Processing: {compound_id}")
        
        # Skip if both CID and SMILES already have values
        if not pd.isna(cid) and not pd.isna(smiles) and cid != "" and smiles != "":
            print(f"  SKIPPED - Already has CID ({cid}) and SMILES")
            rows_skipped += 1
            continue
        
        search_successful = False
        found_cid = None
        
        # Try CAS search first if available
        if not pd.isna(cas) and cas != "":
            print(f"  Searching by CAS: {cas}")
            found_cid, found_smiles = search_by_cas(cas)
            if found_cid and found_smiles:
                df.at[index, "CID号"] = found_cid
                df.at[index, "SMILES"] = found_smiles
                print(f"  SUCCESS via CAS - Found CID: {found_cid}")
                print(f"  SMILES obtained: {found_smiles[:30]}..." if len(found_smiles) > 30 else f"  SMILES obtained: {found_smiles}")
                search_successful = True
            else:
                print(f"  FAILED via CAS - No results found")
        else:
            print("  CAS number not available")
        
        # If CAS search failed or no CAS is available, try English name
        if not search_successful and not pd.isna(english_name) and english_name != "":
            print(f"  Searching by English name: {english_name}")
            
            # Try comprehensive search approach from PubChem_API.py
            result = search_by_comprehensive(english_name)
            
            if result['CID'] and result['SMILES']:
                found_cid = result['CID']
                df.at[index, "CID号"] = result['CID']
                df.at[index, "SMILES"] = result['SMILES']
                print(f"  SUCCESS via comprehensive English name search - Found CID: {result['CID']}")
                print(f"  SMILES obtained: {result['SMILES'][:30]}..." if len(result['SMILES']) > 30 else f"  SMILES obtained: {result['SMILES']}")
                search_successful = True
            else:
                print(f"  FAILED via comprehensive English name search - No results found")
                
                # If all else fails, try one more direct method
                found_cid_direct, found_smiles, found_cas = search_by_name(english_name)
                if found_cid_direct and found_smiles:
                    found_cid = found_cid_direct
                    df.at[index, "CID号"] = found_cid_direct
                    df.at[index, "SMILES"] = found_smiles
                    print(f"  SUCCESS via direct English name search - Found CID: {found_cid_direct}")
                    print(f"  SMILES obtained: {found_smiles[:30]}..." if len(found_smiles) > 30 else f"  SMILES obtained: {found_smiles}")
                    search_successful = True
                else:
                    print(f"  FAILED via all English name methods - No results found")
        else:
            if not search_successful:
                print("  English name not available or previous search successful")
        
        # If we found a CID, extract all CAS numbers from synonyms
        if search_successful and found_cid:
            all_cas_numbers = extract_all_cas_numbers(found_cid)
            if all_cas_numbers:
                # Get existing CAS numbers
                existing_cas = "" if pd.isna(cas) else str(cas)
                existing_cas_list = [c.strip() for c in existing_cas.split(";") if c.strip()]
                
                # Combine existing and new CAS numbers, eliminating duplicates
                combined_cas_list = existing_cas_list.copy()
                new_cas_added = 0
                
                for cas_number in all_cas_numbers:
                    if cas_number not in combined_cas_list:
                        combined_cas_list.append(cas_number)
                        new_cas_added += 1
                
                # Join all CAS numbers with semicolons
                cas_string = ";".join(combined_cas_list)
                df.at[index, "CAS号"] = cas_string
                
                print(f"  Added {new_cas_added} new CAS numbers")
                print(f"  Updated CAS numbers: {cas_string}")
            
            rows_updated += 1
        
        # If all searches failed
        if not search_successful:
            print("  FAILED - Could not find compound information")
            rows_failed += 1
        
        # Save progress every 50 rows in case of interruption
        if index % 50 == 0 and index > 0:
            try:
                df.to_excel(output_file, index=False)
                print(f"\nProgress saved at row {index+1}")
            except Exception as e:
                print(f"\nError saving progress: {e}")
    
    # Save the final updated file
    try:
        df.to_excel(output_file, index=False)
        print("\n" + "="*50)
        print("PROCESSING COMPLETED")
        print("="*50)
        print(f"File saved to: {output_file}")
        print(f"Summary:")
        print(f"  - Rows updated: {rows_updated}")
        print(f"  - Rows skipped: {rows_skipped} (already had CID and SMILES)")
        print(f"  - Rows failed: {rows_failed}")
        print("="*50)
    except Exception as e:
        print(f"\nError saving file: {e}")

def extract_cas_pattern(text):
    """Extract CAS number from text"""
    # Pattern for CAS number: numbers-numbers-number surrounded by whitespace or start/end
    cas_pattern = r'(?:^|\s)(\d+-\d+-\d+)(?:\s|$)'
    match = re.search(cas_pattern, text)
    return match.group(1) if match else ''

def extract_all_cas_numbers(cid):
    """Extract all CAS numbers from a compound's synonyms"""
    try:
        # Get all synonyms for the compound
        synonyms_url = f"https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/cid/{cid}/synonyms/JSON"
        response = requests.get(synonyms_url)
        time.sleep(0.5)  # Add delay
        
        if response.status_code == 200:
            cas_numbers = []
            try:
                synonyms = response.json()['InformationList']['Information'][0]['Synonym']
                
                # Look for all CAS numbers in synonyms
                for synonym in synonyms:
                    # Check if the synonym looks like an isolated CAS number
                    cas_pattern = r'^\s*(\d+-\d+-\d+)\s*$'
                    match = re.search(cas_pattern, synonym)
                    if match:
                        cas_numbers.append(match.group(1))
                        continue
                    
                    # Check if the synonym contains a CAS number surrounded by whitespace
                    cas_pattern = r'(?:^|\s)(\d+-\d+-\d+)(?:\s|$)'
                    for match in re.finditer(cas_pattern, synonym):
                        cas_number = match.group(1)
                        if cas_number not in cas_numbers:
                            cas_numbers.append(cas_number)
                
                if cas_numbers:
                    print(f"  Found {len(cas_numbers)} CAS numbers in synonyms")
                    return cas_numbers
            except (KeyError, IndexError) as e:
                print(f"  Error extracting synonyms: {e}")
                
        return []
    except Exception as e:
        print(f"  Error extracting CAS numbers: {e}")
        return []

def search_by_cas(cas):
    """Search PubChem by CAS number and return CID and SMILES"""
    url = f"https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/name/{cas}/property/IsomericSMILES/JSON"
    try:
        response = requests.get(url)
        time.sleep(0.5)  # Add delay to avoid overloading the API
        if response.status_code == 200:
            data = response.json()
            try:
                cid = data['PropertyTable']['Properties'][0]['CID']
                smiles = data['PropertyTable']['Properties'][0]['IsomericSMILES']
                return cid, smiles
            except (KeyError, IndexError):
                return None, None
        return None, None
    except Exception as e:
        print(f"  Error searching CAS {cas}: {e}")
        return None, None

def search_by_name(name):
    """Search PubChem by English name and return CID, SMILES, and CAS"""
    url = f"https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/name/{name}/property/IsomericSMILES,CAS/JSON"
    try:
        response = requests.get(url)
        time.sleep(0.5)  # Add delay to avoid overloading the API
        if response.status_code == 200:
            data = response.json()
            try:
                cid = data['PropertyTable']['Properties'][0]['CID']
                smiles = data['PropertyTable']['Properties'][0]['IsomericSMILES']
                cas = data['PropertyTable']['Properties'][0].get('CAS', "")  # Use get() to handle missing CAS
                return cid, smiles, cas
            except (KeyError, IndexError):
                return None, None, None
        return None, None, None
    except Exception as e:
        print(f"  Error searching name {name}: {e}")
        return None, None, None

def search_by_comprehensive(compound_name):
    """
    Comprehensive search for a compound using the approach from PubChem_API.py.
    Searches by name, gets CID, then searches for CAS in synonyms, and gets SMILES.
    """
    try:
        print(f"  Trying comprehensive search strategy for: {compound_name}")
        
        # URL encode the compound name
        encoded_name = quote(compound_name)
        
        # First, search for the compound to get its CID
        search_url = f"https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/name/{encoded_name}/cids/JSON"
        response = requests.get(search_url)
        time.sleep(0.5)  # Add delay
        
        if response.status_code == 200:
            try:
                cid = str(response.json()['IdentifierList']['CID'][0])  # Get the best match
                print(f"  Found CID: {cid}")
                
                # Get SMILES using the CID
                smiles_url = f"https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/cid/{cid}/property/IsomericSMILES/JSON"
                smiles_response = requests.get(smiles_url)
                time.sleep(0.5)  # Add delay
                
                if smiles_response.status_code == 200:
                    try:
                        smiles = smiles_response.json()['PropertyTable']['Properties'][0]['IsomericSMILES']
                        print(f"  Found SMILES")
                        
                        # We'll get CAS numbers in the main function using extract_all_cas_numbers
                        return {
                            'CID': cid,
                            'SMILES': smiles,
                            'CAS': ''  # We'll fill this separately
                        }
                    except (KeyError, IndexError) as e:
                        print(f"  Error extracting SMILES: {e}")
            
            except (KeyError, IndexError) as e:
                print(f"  Error extracting CID: {e}")
        
        print(f"  Comprehensive search failed for: {compound_name}")
        return {'CID': '', 'SMILES': '', 'CAS': ''}
    
    except Exception as e:
        print(f"  Error in comprehensive search for {compound_name}: {str(e)}")
        return {'CID': '', 'SMILES': '', 'CAS': ''}

# Example usage
if __name__ == "__main__":
    input_path = "temp_matched_API.xlsx"  
    output_path = "temp_matched_API_updated.xlsx" 
    
    # Test the specific case of "agmatine"
    print("Testing search for 'agmatine'...")
    result = search_by_comprehensive("agmatine")
    if result['CID'] and result['SMILES']:
        print(f"Found agmatine: CID={result['CID']}")
        print(f"SMILES: {result['SMILES']}")
        
        # Extract all CAS numbers
        all_cas = extract_all_cas_numbers(result['CID'])
        if all_cas:
            print(f"Found CAS numbers: {';'.join(all_cas)}")
        else:
            print("No CAS numbers found in synonyms")
    else:
        print("Failed to find agmatine with comprehensive method, trying direct method...")
        cid, smiles, cas = search_by_name("agmatine")
        if cid and smiles:
            print(f"Found agmatine with direct method: CID={cid}, CAS={cas}")
            print(f"SMILES: {smiles}")
            
            # Extract all CAS numbers
            all_cas = extract_all_cas_numbers(cid)
            if all_cas:
                print(f"Found CAS numbers: {';'.join(all_cas)}")
            else:
                print("No CAS numbers found in synonyms")
        else:
            print("Failed to find agmatine with both methods")
    
    print("\nStarting full file processing...")
    search_pubchem(input_path, output_path)